// Copyright (c) 2019-present Dmitry Stepanov and Fyrox Engine contributors.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use crate::{
    asset::manager::ResourceManager,
    core::{
        algebra::{Matrix4, Vector2},
        color::Color,
        math::Rect,
        value_as_u8_slice, ImmutableString,
    },
    graphics::{
        error::FrameworkError,
        framebuffer::{Attachment, DrawCallStatistics, GpuFrameBuffer, ReadTarget},
        gpu_texture::{GpuTexture, GpuTextureDescriptor, GpuTextureKind, PixelKind},
        server::GraphicsServer,
    },
    renderer::{
        bloom::BloomRenderer,
        cache::{
            shader::{binding, property, PropertyGroup, RenderMaterial},
            texture::TextureCache,
            uniform::UniformBufferCache,
        },
        hdr::{adaptation::AdaptationChain, luminance::luminance_evaluator::LuminanceEvaluator},
        make_viewport_matrix,
        resources::RendererResources,
        LuminanceCalculationMethod, QualitySettings, RenderPassStatistics,
    },
    scene::camera::{ColorGradingLut, Exposure},
};

mod adaptation;
mod luminance;

pub struct LumBuffer {
    framebuffer: GpuFrameBuffer,
    size: usize,
}

impl LumBuffer {
    fn new(server: &dyn GraphicsServer, size: usize) -> Result<Self, FrameworkError> {
        let texture =
            server.create_2d_render_target("LuminanceTexture", PixelKind::R32F, size, size)?;
        Ok(Self {
            framebuffer: server.create_frame_buffer(None, vec![Attachment::color(texture)])?,
            size,
        })
    }

    fn clear(&self) {
        self.framebuffer.clear(
            Rect::new(0, 0, self.size as i32, self.size as i32),
            Some(Color::BLACK),
            None,
            None,
        );
    }

    fn matrix(&self) -> Matrix4<f32> {
        make_viewport_matrix(Rect::new(0, 0, self.size as i32, self.size as i32))
    }

    fn texture(&self) -> &GpuTexture {
        &self.framebuffer.color_attachments()[0].texture
    }
}

pub struct HighDynamicRangeRenderer {
    adaptation_chain: AdaptationChain,
    downscale_chain: [LumBuffer; 6],
    frame_luminance: LumBuffer,
    stub_lut: GpuTexture,
    /// Bloom contains only overly bright pixels that create light
    /// bleeding effect (glow effect).
    bloom_renderer: BloomRenderer,
}

pub struct HdrRendererArgs<'a> {
    pub server: &'a dyn GraphicsServer,
    pub hdr_scene_frame: &'a GpuTexture,
    pub ldr_framebuffer: &'a GpuFrameBuffer,
    pub viewport: Rect<i32>,
    pub dt: f32,
    pub exposure: Exposure,
    pub color_grading_lut: Option<&'a ColorGradingLut>,
    pub use_color_grading: bool,
    pub texture_cache: &'a mut TextureCache,
    pub uniform_buffer_cache: &'a mut UniformBufferCache,
    pub renderer_resources: &'a RendererResources,
    pub resource_manager: &'a ResourceManager,
    pub settings: &'a QualitySettings,
}

impl HighDynamicRangeRenderer {
    pub fn new(
        width: usize,
        height: usize,
        server: &dyn GraphicsServer,
    ) -> Result<Self, FrameworkError> {
        Ok(Self {
            frame_luminance: LumBuffer::new(server, 64)?,
            downscale_chain: [
                LumBuffer::new(server, 32)?,
                LumBuffer::new(server, 16)?,
                LumBuffer::new(server, 8)?,
                LumBuffer::new(server, 4)?,
                LumBuffer::new(server, 2)?,
                LumBuffer::new(server, 1)?,
            ],
            adaptation_chain: AdaptationChain::new(server)?,
            stub_lut: server.create_texture(GpuTextureDescriptor {
                name: "StubHdrLut",
                kind: GpuTextureKind::Volume {
                    width: 1,
                    height: 1,
                    depth: 1,
                },
                pixel_kind: PixelKind::RGB8,
                data: Some(&[0, 0, 0]),
                ..Default::default()
            })?,
            bloom_renderer: BloomRenderer::new(server, width, height)?,
        })
    }

    fn calculate_frame_luminance(
        &self,
        server: &dyn GraphicsServer,
        scene_frame: &GpuTexture,
        uniform_buffer_cache: &mut UniformBufferCache,
        renderer_resources: &RendererResources,
    ) -> Result<DrawCallStatistics, FrameworkError> {
        let _debug_scope = server.begin_scope("CalculateFrameLuminance");

        self.frame_luminance.clear();

        let frame_matrix = self.frame_luminance.matrix();
        let inv_size = Vector2::repeat(1.0 / self.frame_luminance.size as f32);

        let properties = PropertyGroup::from([
            property("worldViewProjection", &frame_matrix),
            property("invSize", &inv_size),
        ]);
        let material = RenderMaterial::from([
            binding(
                "frameSampler",
                (scene_frame, &renderer_resources.nearest_clamp_sampler),
            ),
            binding("properties", &properties),
        ]);

        renderer_resources.shaders.hdr_luminance.run_pass(
            1,
            &ImmutableString::new("Primary"),
            &self.frame_luminance.framebuffer,
            &renderer_resources.quad,
            Rect::new(
                0,
                0,
                self.frame_luminance.size as i32,
                self.frame_luminance.size as i32,
            ),
            &material,
            uniform_buffer_cache,
            Default::default(),
            None,
        )
    }

    fn calculate_avg_frame_luminance(
        &self,
        server: &dyn GraphicsServer,
        uniform_buffer_cache: &mut UniformBufferCache,
        renderer_resources: &RendererResources,
        luminance_calculation_method: LuminanceCalculationMethod,
    ) -> Result<RenderPassStatistics, FrameworkError> {
        let _debug_scope = server.begin_scope("CalculateAvgFrameLuminance");

        let mut stats = RenderPassStatistics::default();

        match luminance_calculation_method {
            LuminanceCalculationMethod::Histogram => {
                // TODO: Cloning memory from GPU to CPU is slow, but since the engine is limited
                // by macOS's OpenGL 4.1 support and lack of compute shaders we'll build histogram
                // manually on CPU anyway. Replace this with compute shaders whenever possible.
                let pixels = self
                    .frame_luminance
                    .framebuffer
                    .read_pixels_of_type::<f32>(ReadTarget::Color(0))
                    .ok_or_else(|| {
                        FrameworkError::Custom("Unable to read luminance buffer!".to_string())
                    })?;

                let evaluator =
                    luminance::histogram_luminance_evaluator::HistogramLuminanceEvaluator::default(
                    );
                let avg_value = evaluator.average_luminance(&pixels);

                self.downscale_chain.last().unwrap().texture().set_data(
                    GpuTextureKind::Rectangle {
                        width: 1,
                        height: 1,
                    },
                    PixelKind::R32F,
                    1,
                    Some(value_as_u8_slice(&avg_value)),
                )?;
            }
            LuminanceCalculationMethod::DownSampling => {
                let mut src = &self.frame_luminance;

                for dest in self.downscale_chain.iter() {
                    let src_inv_size = Vector2::repeat(1.0 / src.size as f32);
                    let src_texture = src.texture();

                    let matrix = dest.matrix();

                    let properties = PropertyGroup::from([
                        property("worldViewProjection", &matrix),
                        property("invSize", &src_inv_size),
                    ]);
                    let material = RenderMaterial::from([
                        binding(
                            "lumSampler",
                            (src_texture, &renderer_resources.linear_clamp_sampler),
                        ),
                        binding("properties", &properties),
                    ]);

                    stats += renderer_resources.shaders.hdr_downscale.run_pass(
                        1,
                        &ImmutableString::new("Primary"),
                        &dest.framebuffer,
                        &renderer_resources.quad,
                        Rect::new(0, 0, dest.size as i32, dest.size as i32),
                        &material,
                        uniform_buffer_cache,
                        Default::default(),
                        None,
                    )?;

                    src = dest;
                }
            }
        }

        Ok(stats)
    }

    fn adaptation(
        &self,
        server: &dyn GraphicsServer,
        dt: f32,
        uniform_buffer_cache: &mut UniformBufferCache,
        renderer_resources: &RendererResources,
    ) -> Result<DrawCallStatistics, FrameworkError> {
        let _debug_scope = server.begin_scope("Adaptation");

        let ctx = self.adaptation_chain.begin();
        let viewport = Rect::new(0, 0, ctx.lum_buffer.size as i32, ctx.lum_buffer.size as i32);
        let matrix = ctx.lum_buffer.matrix();

        let speed = 0.3 * dt;
        let properties = PropertyGroup::from([
            property("worldViewProjection", &matrix),
            property("speed", &speed),
        ]);
        let material = RenderMaterial::from([
            binding(
                "oldLumSampler",
                (&ctx.prev_lum, &renderer_resources.nearest_clamp_sampler),
            ),
            binding(
                "newLumSampler",
                (
                    self.downscale_chain.last().unwrap().texture(),
                    &renderer_resources.nearest_clamp_sampler,
                ),
            ),
            binding("properties", &properties),
        ]);

        renderer_resources.shaders.hdr_adaptation.run_pass(
            1,
            &ImmutableString::new("Primary"),
            &ctx.lum_buffer.framebuffer,
            &renderer_resources.quad,
            viewport,
            &material,
            uniform_buffer_cache,
            Default::default(),
            None,
        )
    }

    fn map_hdr_to_ldr(&self, args: HdrRendererArgs) -> Result<DrawCallStatistics, FrameworkError> {
        let HdrRendererArgs {
            server,
            hdr_scene_frame,
            ldr_framebuffer,
            viewport,
            exposure,
            color_grading_lut,
            use_color_grading,
            texture_cache,
            uniform_buffer_cache,
            renderer_resources,
            resource_manager,
            settings,
            ..
        } = args;

        let _debug_scope = args.server.begin_scope("ToneMap");

        let frame_matrix = make_viewport_matrix(viewport);

        let color_grading_lut_tex = color_grading_lut
            .and_then(|l| {
                texture_cache
                    .get(server, resource_manager, l.lut_ref())
                    .map(|t| (&t.gpu_texture, &t.gpu_sampler))
            })
            .unwrap_or((&self.stub_lut, &renderer_resources.nearest_clamp_sampler));

        let (is_auto, min_luminance, max_luminance, fixed_exposure) = match exposure {
            Exposure::Auto {
                min_luminance,
                max_luminance,
            } => (true, min_luminance, max_luminance, 0.0),
            Exposure::Manual(fixed_exposure) => (false, 0.0, 0.0, fixed_exposure),
        };

        let bloom_texture = if settings.hdr_settings.bloom_settings.use_bloom {
            self.bloom_renderer.result()
        } else {
            &renderer_resources.black_dummy
        };

        let color_grading_enabled = use_color_grading && color_grading_lut.is_some();
        let properties = PropertyGroup::from([
            property("worldViewProjection", &frame_matrix),
            property("useColorGrading", &color_grading_enabled),
            property("minLuminance", &min_luminance),
            property("maxLuminance", &max_luminance),
            property("autoExposure", &is_auto),
            property("fixedExposure", &fixed_exposure),
        ]);
        let material = RenderMaterial::from([
            binding(
                "hdrSampler",
                (hdr_scene_frame, &renderer_resources.nearest_clamp_sampler),
            ),
            binding(
                "lumSampler",
                (
                    self.adaptation_chain.avg_lum_texture(),
                    &renderer_resources.nearest_clamp_sampler,
                ),
            ),
            binding(
                "bloomSampler",
                (bloom_texture, &renderer_resources.linear_clamp_sampler),
            ),
            binding("colorMapSampler", color_grading_lut_tex),
            binding("properties", &properties),
        ]);

        renderer_resources.shaders.hdr_map.run_pass(
            1,
            &ImmutableString::new("Primary"),
            ldr_framebuffer,
            &renderer_resources.quad,
            viewport,
            &material,
            uniform_buffer_cache,
            Default::default(),
            None,
        )
    }

    pub fn render(&self, args: HdrRendererArgs) -> Result<RenderPassStatistics, FrameworkError> {
        let _debug_scope = args.server.begin_scope("HDR");
        let mut stats = RenderPassStatistics::default();
        stats += self.calculate_frame_luminance(
            args.server,
            args.hdr_scene_frame,
            args.uniform_buffer_cache,
            args.renderer_resources,
        )?;
        stats += self.calculate_avg_frame_luminance(
            args.server,
            args.uniform_buffer_cache,
            args.renderer_resources,
            args.settings.hdr_settings.luminance_calculation_method,
        )?;
        if args.settings.hdr_settings.bloom_settings.use_bloom {
            stats += self.bloom_renderer.render(
                args.server,
                args.hdr_scene_frame,
                args.uniform_buffer_cache,
                args.renderer_resources,
                args.settings,
            )?;
        }
        stats += self.adaptation(
            args.server,
            args.dt,
            args.uniform_buffer_cache,
            args.renderer_resources,
        )?;
        stats += self.map_hdr_to_ldr(args)?;
        Ok(stats)
    }
}
